#pragma once

#include <ocs2_core/dynamics/SystemDynamicsBase.h>

#include "mpc_quadrotor_mujoco/definitions.h"
#include "mpc_quadrotor_mujoco/QuadrotorParameters.h"

namespace ocs2 {
    namespace quadrotor {

        template <typename SCALAR_T>
        Eigen::Matrix<SCALAR_T, 3, 3> getMppingFromLocalAngularVelocityToRPYDerivative(const Eigen::Matrix<SCALAR_T, 3, 1>& rpyAngles) {
            const SCALAR_T sx = sin(rpyAngles(0));
            const SCALAR_T cx = cos(rpyAngles(0));
            const SCALAR_T cy = cos(rpyAngles(1));
            const SCALAR_T ty = tan(rpyAngles(1));

            Eigen::Matrix<SCALAR_T, 3, 3> M;
            M << SCALAR_T(1.0), ty * sx, ty * cx,
                SCALAR_T(0.0), cx, -sx,
                SCALAR_T(0.0), sx / cy, cx / cy;
            return M;
        }

        template <typename SCALAR_T>
        inline Eigen::Matrix<SCALAR_T, 3, 6> JacobianOfAngularVelocityMapping(const Eigen::Matrix<SCALAR_T, 3, 1>& rpyAngles, const Eigen::Matrix<SCALAR_T, 3, 1>& angularVeolocity) {
            const SCALAR_T sx = sin(rpyAngles(0));
            const SCALAR_T cx = cos(rpyAngles(0));
            const SCALAR_T sy = sin(rpyAngles(1));
            const SCALAR_T cy = cos(rpyAngles(1));
            const SCALAR_T ty = tan(rpyAngles(1));

            Eigen::Matrix<SCALAR_T, 3, 6> jac;
            jac.block(0, 3, 3, 3) = getMppingFromLocalAngularVelocityToRPYDerivative(rpyAngles);
            
            jac.block(0, 0, 3, 3).setZero();
            jac(0, 0) = (cx * angularVeolocity(1) - sx * angularVeolocity(2)) * ty;
            jac(0, 1) = (sx * angularVeolocity(1) + cx * angularVeolocity(2)) / (cy * cy);

            jac(1, 0) = -sx * angularVeolocity(1) - cx * angularVeolocity(2);

            jac(2, 0) = (cx * angularVeolocity(1) - sx * angularVeolocity(2)) / cy;
            jac(2, 1) = (sx * angularVeolocity(1) + cx * angularVeolocity(2)) * (ty / cy);
            
            return jac;
        }

        class QuadrotorSystemDynamics final : public SystemDynamicsBase {
        public:
            explicit QuadrotorSystemDynamics(QuadrotorParameters quadrotorParameters) : param_(std::move(quadrotorParameters)) {}

            ~QuadrotorSystemDynamics() override = default;

            QuadrotorSystemDynamics* clone() const override { return new QuadrotorSystemDynamics(*this); }

            vector_t computeFlowMap(scalar_t time, const vector_t& state, const vector_t& input, const PreComputation&) override;
            VectorFunctionLinearApproximation linearApproximation(scalar_t t, const vector_t& x, const vector_t& u, const PreComputation& preComp) override;

        private:
            QuadrotorParameters param_;

            matrix_t jacobianOfAngularVelocityMapping_;
        };
    } // namespace quadrotor
} // namespace ocs2
